import os
import numpy as np
import rasterio
from rasterio.mask import mask
import geopandas as gpd
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
from matplotlib import rcParams

# 设置中文字体
rcParams['font.sans-serif'] = ['Noto Sans CJK JP']  # 选择简体中文
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 文件路径
rainfall_tif = './mean_rain_tem_outputs/mean_rainfall.tif'  # 降雨平均值文件
temperature_tif = './mean_rain_tem_outputs/mean_temperature.tif'  # 气温平均值文件
mask_shp = './shanxi_mask/shanxi.shp'  # 掩膜文件

# 加载掩膜数据
mask_gdf = gpd.read_file(mask_shp)

# 定义绘图函数
def plot_tif_with_mask(tif_file, mask_gdf, ax, title, cmap, cbar_label):
    with rasterio.open(tif_file) as src:
        # 使用掩膜裁剪数据
        out_image, out_transform = mask(src, mask_gdf.geometry, crop=True)
        data = out_image[0]  # 读取第一个波段
        # data = data/10
        nodata = src.nodata
        if nodata is not None:
            data = np.where(data == nodata, np.nan, data)  # 将 NoData 值替换为 NaN

        # 绘制数据
        img = ax.imshow(data, cmap=cmap, extent=[
            out_transform[2],  # 左
            out_transform[2] + out_transform[0] * data.shape[1],  # 右
            out_transform[5] + out_transform[4] * data.shape[0],  # 下
            out_transform[5]  # 上
        ], origin='upper', norm=Normalize(vmin=np.nanmin(data), vmax=np.nanmax(data)))

        # 添加掩膜边界
        mask_gdf.boundary.plot(ax=ax, edgecolor='black', linewidth=1)

        # 添加标题
        ax.set_title(title, fontsize=14)

        # 添加颜色条
        cbar = plt.colorbar(img, ax=ax, orientation='vertical', shrink=0.7)
        cbar.set_label(cbar_label, fontsize=12)

# 创建绘图
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 8))

# 绘制降雨图
plot_tif_with_mask(
    rainfall_tif, mask_gdf, ax1,
    title='陕西省年降雨量/mm',
    cmap='YlGnBu',  # 蓝绿色渐变
    cbar_label='降雨量 (mm)'
)

# 绘制气温图
plot_tif_with_mask(
    temperature_tif, mask_gdf, ax2,
    title='陕西省年平均气温/°C',
    cmap='RdYlBu_r',  # 红蓝渐变
    cbar_label='气温 (°C)'
)

# 调整布局
plt.tight_layout()
plt.savefig('rainfall_temperature_map.png', dpi=300)
plt.show()